﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace GpuProgramCompiler.Compilers
{
    abstract class ShaderCompilerBase
    {
        private ShaderX srcShaderX;
		private CmdLineArgs cmdLineArgs;
		private CompiledShader compiledShader;

        protected class ReplacementTerm
        {
            public String text;
            public int start;
            public int end;
        }

        protected ShaderX SrcShaderX
        {
            get { return srcShaderX; }
        }

        protected CmdLineArgs Args
        {
            get { return cmdLineArgs; }
        }

        protected CompiledShader DstCompiledShader
        {
            get { return compiledShader; }
        }

        protected abstract void ExpandShader(ShaderX.CodeBlock codeBlock, String type);

        public void Compile(CmdLineArgs args, ShaderX shaderX, ref CompiledShader outCompiledShader)
		{
			srcShaderX = shaderX;
			cmdLineArgs = args;
			compiledShader = outCompiledShader;

			foreach (ShaderX.Shader shader in shaderX.Shaders)
			{
				foreach (ShaderX.ShaderStage stage in shader.stages)
				{
					ExpandShader(shaderX.FindCodeBlock(stage.code), stage.type);
				}
			}
		}

        protected ReplacementTerm ExtractReplacementTerm(String data)
        {
            int startOfPattern = data.IndexOf("$");
            if (startOfPattern != -1)
            {
                int endOfPattern = startOfPattern;
                int indent = 0;
                for (int i = startOfPattern; i < data.Length; ++i)
                {
                    if (data[i] == '(')
                        indent++;
                    else if (data[i] == ')')
                    {
                        if (--indent == 0)
                        {
                            endOfPattern = i + 1;
                            break;
                        }
                    }
                }

                if (endOfPattern > startOfPattern)
                {
                    ReplacementTerm rt = new ReplacementTerm();
                    rt.start = startOfPattern;
                    rt.end = endOfPattern;
                    rt.text = data.Substring(startOfPattern + 2, endOfPattern - startOfPattern - 3);
                    return rt;
                }
            }

            return null;
        }

        protected void ReplaceTerm(ref String code, ReplacementTerm replacementTerm, String replacement)
        {
            code = code.Remove(replacementTerm.start, replacementTerm.end - replacementTerm.start);
            code = code.Insert(replacementTerm.start, replacement);
        }

        protected String DoVariableReplacement(String code, Dictionary<String, String> variables)
        {
            foreach (String variable in variables.Keys)
            {
                String varName = "$(" + variable + ")";
                code = code.Replace(varName, variables[variable]);
            }

            return code;
        }

        protected String ExpandFunction(ShaderX.CodeBlock codeBlock, List<ShaderX.IOElement> inputs, List<ShaderX.IOElement> outputs, List<ShaderX.IOElement> uniforms, Dictionary<String, String> variables)
        {
            // First process inputs and outputs for this block
            foreach (ShaderX.IOElement ioElement in codeBlock.ioElements)
            {
                if (ioElement.io == "Input")
                {
                    String inputPrefix = "";
                    if (ioElement.source == "external")
                    {
                        inputs.Add(ioElement);
                        inputPrefix = "INPUT.";
                    }
                    else if (ioElement.source == "uniform")
                    {
                        uniforms.Add(ioElement);
                        variables[ioElement.name] = ioElement.name;
                    }

                    if (!variables.ContainsKey(ioElement.name))
                        variables[ioElement.name] = inputPrefix + ioElement.name;
                }
                else if (ioElement.io == "Output")
                {
                    if (ioElement.source == "external")
                    {
                        outputs.Add(ioElement);
                    }

                    if (!variables.ContainsKey(ioElement.name))
                        variables[ioElement.name] = "OUTPUT." + ioElement.name;
                }
            }

            // Process code for this block
            String code = codeBlock.code;

            code = DoVariableReplacement(code, variables);

            // Now go through functions
            ReplacementTerm replacementTerm = ExtractReplacementTerm(code);
            while (replacementTerm != null)
            {
                String[] terms = replacementTerm.text.Split(',');
                if (terms.Length > 0)
                {
                    ShaderX.CodeBlock functionCodeBlock = SrcShaderX.FindCodeBlock(terms[0]);
                    if (functionCodeBlock == null)
                    {
                        throw new Exception(String.Format("Unknown symbol '{0}' found in code block '{1}'", terms[0], codeBlock.name));
                    }

                    Dictionary<String, String> functionVariables = new Dictionary<String, String>();
                    for (int i = 1; i < terms.Length; ++i)
                    {
                        String[] paramBits = terms[i].Split('=');
                        if (paramBits.Length == 2)
                        {
                            functionVariables[paramBits[0]] = paramBits[1];
                        }
                        else
                        {
                            throw new Exception(String.Format("Incorrect function variable assignment '{0}' in code block '{1}'", replacementTerm.text, codeBlock.name));
                        }
                    }
                    String functionCode = ExpandFunction(functionCodeBlock, inputs, outputs, uniforms, functionVariables);
                    ReplaceTerm(ref code, replacementTerm, functionCode);
                }
                else
                {
                    throw new Exception(String.Format("Dubious use of $ syntax in code '{0}' in code block '{1}'", replacementTerm.text, codeBlock.name));
                }

                replacementTerm = ExtractReplacementTerm(code);
            }

            String[] lines = code.Split(';');
            code = "";
            foreach (String line in lines)
            {
                String newLine = "\t" + line.Trim() + ";\n";
                if (newLine.Trim() != ";")
                {
                    code += newLine;
                }
            }

            return code;
        }
    }
}
